import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import torch
import cv2
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
from func import load,product,preprocess
from network import ss_cnn, operate

## 指定显卡

USE_GPU=True
if USE_GPU:
    os.environ["CUDA_VISIBLE_DEVICES"] = "3"
else:
    device=torch.device('cpu')

a=load()

All_data, _, _, _,r,c,FLAG=a.load_data(flag='ksc')
###################### 加载数据, 模型及设置参数及全局变量 ####################
# # 设置归一化范围（因加载image_3d_mat_origin，inference时无需再次归一化）
# mi = -0.5
# ma = 0.5
# Patch块尺寸-1的一半,手动指定
half_s=13
#每块样本数
BATCH=50000

trn_num=np.load('trn_num_'+str(FLAG)+'.npy')
pre_num=np.load('pre_num_'+str(FLAG)+'.npy')
y_trn=np.load('y_trn_'+str(FLAG)+'.npy')
image_3d_mat_origin=np.load('image_3d_mat_origin_'+str(FLAG)+'.npy')

###### 网络相关情况设置
net=torch.load('3dcnn_'+str(FLAG)+'.pkl',map_location='cpu')
#net=net.module#if use DataParallel
net=net.cuda()

criterion = torch.nn.NLLLoss()  # 负对数似然损失函数（如果不算log_softmax则直接采用交叉熵损失函数)

###### 标签相关情况设置
y_disp=np.zeros([All_data.shape[0]])
y_disp[trn_num]=y_trn

y_disp_all=y_disp.copy()

start=0
end=np.min([start+BATCH,pre_num.shape[0]])

###################################### 预测 #######################################

part_num=int(pre_num.shape[0]/BATCH)+1

print('需要分成{}块来预测'.format(part_num))

for i in range(0,part_num):

    ###################### label ######################

    pre_num_part=pre_num[start:end]

    y_pre=All_data[pre_num_part,-1]#include background

    pre_YY = torch.from_numpy(np.ones([y_pre.shape[0]]))

    ######################  data ######################

    a = product(c, FLAG)

    pre_spat, pre_num_part = a.production_data_valtespre(pre_num_part, half_s, image_3d_mat_origin, flag='Pre')

    pre_spat = pre_spat[:, :, :, :, np.newaxis]

    pre_XX_spat=torch.from_numpy(pre_spat.transpose(0, 4, 3, 1, 2))

    del pre_spat


    ######### 推断，预测集 #########

    pre_dataset=TensorDataset(pre_XX_spat,pre_YY)
    pre_loader=DataLoader(pre_dataset,batch_size=300)

    #net=nn.DataParallel(net,device_ids=[0])
    #net=net.cpu()

    a=operate()
    y_pred_pre=a.inference(net,pre_loader,criterion,FLAG='PRED')

    print('第{}块数据预测完成！！！'.format(i))

    y_disp_all[pre_num_part]=y_pred_pre

    start=end
    end=np.min([start+BATCH,pre_num.shape[0]])

######################  展示   ##########################

#np.save('y_pre_before10w.npy',y_disp_all)

# plt.xlabel('All image')
# plt.imshow(y_disp_all.reshape(r,c),cmap='jet')
# plt.xticks([])
# plt.yticks([])
#
# plt.show()

#####################  保存预测全图 #####################

#保存生成的预测图便于后期上色

cv2.imwrite('3dcnn_all_'+str(FLAG)+'.png', y_disp_all.reshape(r,c))

# plt.subplots(figsize=[10,10])
# a1=plt.imshow(y_disp_all.reshape(r,c),cmap='jet')
# plt.xticks([])
# plt.yticks([])
# plt.savefig('3dcnn_all_'+str(flag)+'.png',dpi=600,bbox_inches='tight')
# plt.show()

print('结果图展示 & 保存阶段完成！！')